#ifndef YOLO_DETECTOR_H
#define YOLO_DETECTOR_H

#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <string>
#include <vector>

/**
 * @brief 检测结果的结构体，包含边界框、类别和置信度信息
 */
struct Detection {
    int class_id;          // 类别ID
    float confidence;      // 置信度
    cv::Rect box;          // 边界框
    std::string className; // 类别名称
};

/**
 * @brief YOLOv8目标检测器类
 * 
 * 封装了基于OpenCV DNN的YOLOv8目标检测功能，支持ONNX模型格式
 * 主要功能包括模型初始化、图像预处理、目标检测以及结果可视化
 */
class YOLODetector {
public:
    /**
     * @brief 构造函数，初始化基本参数
     */
    YOLODetector();
    
    /**
     * @brief 初始化模型
     * 
     * @param modelPath ONNX模型文件的路径
     * @return bool 初始化是否成功
     */
    bool initialize(const std::string& modelPath);
    
    /**
     * @brief 执行目标检测
     * 
     * @param image 输入图像
     * @param confThreshold 置信度阈值，默认0.25
     * @param nmsThreshold NMS阈值，默认0.45
     * @return std::vector<Detection> 检测结果列表
     */
    std::vector<Detection> detect(const cv::Mat& image, float confThreshold = 0.25f, float nmsThreshold = 0.45f);
    
    /**
     * @brief 在图像上绘制检测结果
     * 
     * @param image 输入/输出图像
     * @param detections 检测结果列表
     */
    void drawDetections(cv::Mat& image, const std::vector<Detection>& detections);

private:
    cv::dnn::Net net;               // OpenCV DNN模型
    cv::Size inputSize;             // 输入图像大小
    std::vector<std::string> classNames; // 类别名称列表
    
    /**
     * @brief 图像预处理
     * 
     * @param image 输入图像
     * @return cv::Mat 处理后的图像Blob
     */
    cv::Mat preprocess(const cv::Mat& image);
    
    /**
     * @brief 后处理网络输出
     * 
     * @param output 网络原始输出
     * @param originalImageSize 原始图像尺寸
     * @param confThreshold 置信度阈值
     * @param nmsThreshold NMS阈值
     * @return std::vector<Detection> 检测结果列表
     */
    std::vector<Detection> postprocess(const cv::Mat& output, const cv::Size& originalImageSize, 
                                      float confThreshold, float nmsThreshold);
};

#endif // YOLO_DETECTOR_H 